#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/1/20 17:24
@desc: 
"""
import argparse
import logging
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from crowdcounting.data.tools import get_density_map
from crowdcounting.test.tools import get_image_list
from utils.basic import get_file_name
from utils.io import read_text_file


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--image-root', type=str, default='', dest='root',
        help='The directory where the image data is stored.')
    parser.add_argument(
        '--prefix', help='postfix to file', default='mtl',
        dest='prefix', type=str)
    parser.add_argument(
        '--model-dir', type=str, default='', dest='model',
        help='The directory where the pd model will be stored.')
    parser.add_argument(
        '--min', type=int, default=0, dest='min',
        help='The minimum number of fisheye in per image which would be predict.')
    parser.add_argument(
        '--max', type=int, default=10000, dest='max',
        help='The maximum number of fisheye in per image which would be predict.')
    parser.add_argument(
        '--output-root', type=str, default='', dest='output',
        help='the root for all the output.')
    return parser.parse_args()


def predict_image(image_list, model, size=256):
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    shape = (size, size)
    mae, mse = 0.0, 0.0
    gt_total, pre_total = 0.0, 0.0

    with tf.Session(config=config) as sess:
        meta_graph_def = tf.saved_model.loader.load(sess, ['frozen_model'], model)
        signature = meta_graph_def.signature_def
        image_tensor_name = signature['predict_images'].inputs['images'].name
        pre_density_map_name = signature['predict_images'].outputs['density_map'].name
        image_tensor = sess.graph.get_tensor_by_name(image_tensor_name)
        density_map_name = sess.graph.get_tensor_by_name(pre_density_map_name)
        for idx, paths in enumerate(image_list):
            logging.info('predicting {} of {}'.format(idx, len(image_list)))
            img_path, txt_path = paths.split('&!&')
            img = cv2.imread(img_path)
            img = cv2.resize(img, shape)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img.astype(np.float32)
            img[:, :, 0] -= 80.197326
            img[:, :, 1] -= 75.26667
            img[:, :, 2] -= 72.822747
            gt_number = len(read_text_file(txt_path.strip()))
            img_np = np.expand_dims(img, axis=0)
            density_map = sess.run([density_map_name], feed_dict={image_tensor: img_np})
            pre_number = np.sum(density_map)

            logging.info('gt_number: {}, pre_number: {}'.format(gt_number, pre_number))
            gt_total += gt_number
            pre_total += pre_number
            mae += abs(gt_number - pre_number)
            mse += (gt_number - pre_number) ** 2

    logging.info('MAE: {:.5f}, MSE: {:.5f}'.format(mae / len(image_list), mse / len(image_list)))
    logging.info('AVE: {:.5f}'.format(gt_total / len(image_list)))
    logging.info('accuracy: {:.5f} ----- 1 - MAE/AVE'.format(1.0 - mae / gt_total))
    logging.info('gt_total: {}, pre_total: {:.5f}'.format(gt_total, pre_total))
    logging.info('recall: {:.5f} ----- 1 - abs(gt_total - pre_total) / gt_total'.format(
        1 - abs(gt_total - pre_total) / gt_total))


def predict_save_image(image_list, model, output, size=256):
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    shape = (size, size)

    with tf.Session(config=config) as sess:
        meta_graph_def = tf.saved_model.loader.load(sess, ['frozen_model'], model)
        signature = meta_graph_def.signature_def
        image_tensor_name = signature['predict_images'].inputs['images'].name
        pre_density_map_name = signature['predict_images'].outputs['density_map'].name
        classify_result_name = signature['predict_images'].outputs['class'].name
        for idx, paths in enumerate(image_list):
            image_tensor = sess.graph.get_tensor_by_name(image_tensor_name)
            pre_density_map = sess.graph.get_tensor_by_name(pre_density_map_name)
            classify_result = sess.graph.get_tensor_by_name(classify_result_name)
            logging.info('predicting {} of {}'.format(idx, len(image_list)))
            txt_path, img_path = paths.split('&!&')
            img = cv2.imread(img_path)
            img = cv2.resize(img, shape)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img.astype(np.float32)
            img[:, :, 0] -= 80.197326
            img[:, :, 1] -= 75.26667
            img[:, :, 2] -= 72.822747
            img_np = np.expand_dims(img, axis=0)
            pre_density_map, classify_result = sess.run(
                [pre_density_map, classify_result],
                feed_dict={
                    image_tensor: img_np
                })
            points = read_text_file(txt_path)
            points = list(map(lambda x: x.split(' '), points))
            points = list(map(lambda x: [float(x[0]), float(x[1])], points))
            points_array = np.asarray(points) * size
            points_array = points_array.round().astype(np.int32)
            density_map = get_density_map(shape, points_array, kernel_size=15, sigma=4)
            bins = 88 / 24
            cls = int((np.sum(density_map) - 24) / bins)
            logging.info('predicting class: {}, number: {}'.format(np.argmax(classify_result), np.sum(pre_density_map)))
            logging.info('ground true class: {}, number: {}'.format(cls, np.sum(density_map)))
            plt.figure(figsize=(22, 10), dpi=120)
            plt.subplot(122)
            plt.title('predict class: {}, number: {}'.format(np.argmax(classify_result), np.sum(pre_density_map)))
            plt.imshow(np.squeeze(pre_density_map), interpolation='none')
            plt.subplot(121)
            plt.title('ground class: {}, number: {}'.format(cls, np.sum(density_map)))
            plt.imshow(density_map, interpolation='none')
            plt.savefig(os.path.join(output, '{}.jpg'.format(get_file_name(img_path))))
            plt.close()


if __name__ == '__main__':
    ARGS = parse_args()
    logging.basicConfig(filename=os.path.join(ARGS.output,
                                              '{}_{}_{}.log'.format(ARGS.prefix, ARGS.min, ARGS.max)),
                        level=logging.INFO)
    images = get_image_list(ARGS.root, ARGS.min, ARGS.max)
    # predict_image(images, ARGS.model)
    predict_save_image(images, ARGS.model, ARGS.output)
